from pathlib import Path

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from loguru import logger


class FileEncryptDecryptUtils:
    # 示例密钥（16进制字符串格式）
    EXAMPLE_KEY_256_STRING = (
        "6DA64CC3DA86607C1AF4B193F51B5036E29AB3E35FD6F6375D23ABDA48BC84B8"
    )
    BLOCK_SIZE = 16  # AES块大小为16字节

    @staticmethod
    def hex_string_to_bytes(hex_string):
        return bytes.fromhex(hex_string)

    def encrypt_file_aes256_ecb(
        self, key_hex_string, plain_filename, encrypted_filename
    ):
        print(f"原始文件名: {plain_filename}")
        print(f"输出加密文件名: {encrypted_filename}")

        key_bytes = self.hex_string_to_bytes(key_hex_string)
        cipher = AES.new(key_bytes, AES.MODE_ECB)

        # 读取和加密文件
        with open(plain_filename, "rb") as f_in:
            data = f_in.read()
            # 使用PKCS5Padding进行填充
            padded_data = pad(data, self.BLOCK_SIZE)
            encrypted_data = cipher.encrypt(padded_data)

        with open(encrypted_filename, "wb") as f_out:
            f_out.write(encrypted_data)

    def decrypt_file_aes256_ecb(
        self, key_hex_string, encrypted_filename, decrypted_filename
    ):
        print(f"测试加密文件名: {encrypted_filename}")
        print(f"输出解密文件名: {decrypted_filename}")
        print(f"测试密钥hex字符串: {key_hex_string}")

        key_bytes = self.hex_string_to_bytes(key_hex_string)
        cipher = AES.new(key_bytes, AES.MODE_ECB)

        # 读取和解密文件
        with open(encrypted_filename, "rb") as f_in:
            encrypted_data = f_in.read()
            decrypted_data = cipher.decrypt(encrypted_data)
            unpadded_data = unpad(decrypted_data, self.BLOCK_SIZE)

        with open(decrypted_filename, "wb") as f_out:
            f_out.write(unpadded_data)

    def encrypt(self, input_path, output_path=None):
        cipher = AES.new(
            self.hex_string_to_bytes(self.EXAMPLE_KEY_256_STRING), AES.MODE_ECB
        )

        # 读取文件 使用PKCS5Padding进行填充
        in_name = Path(input_path)
        logger.info(f"加密前文件：{in_name.as_posix()}")
        encrypted_data = cipher.encrypt(pad(in_name.read_bytes(), self.BLOCK_SIZE))

        if output_path is None:
            return encrypted_data

        # 写入文件
        out_name = Path(output_path)
        logger.info(f"加密后文件：{out_name.as_posix()}")
        out_name.write_bytes(encrypted_data)

    def decrypt(self, input_path, output_path=None):
        cipher = AES.new(
            self.hex_string_to_bytes(self.EXAMPLE_KEY_256_STRING), AES.MODE_ECB
        )

        # 读取和解密文件
        in_name = Path(input_path)
        logger.info(f"解密前文件：{in_name.as_posix()}")
        decrypted_data = unpad(cipher.decrypt(in_name.read_bytes()), self.BLOCK_SIZE)

        if output_path is None:
            return decrypted_data

        # 写入文件
        out_name = Path(output_path)
        logger.info(f"解密后文件：{out_name.as_posix()}")
        out_name.write_bytes(decrypted_data)


def demo():
    # 测试明文文件
    cls = FileEncryptDecryptUtils()
    cls.encrypt(
        r"", r"C:\Users\wk\IdeaProjects\untitled\encrypted_client_to_upload.txt"
    )


if __name__ == "__main__":
    demo()
